/*
 * AIEngine a new generation network intrusion detection system.
 *
 * Copyright (C) 2013-2023  Luis Campo Giralte
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Library General Public
 * License as published by the Free Software Foundation; either
 * version 2 of the License, or (at your option) any later version.
 *
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Library General Public License for more details.
 *
 * You should have received a copy of the GNU Library General Public
 * License along with this library; if not, write to the
 * Free Software Foundation, Inc., 51 Franklin St, Fifth Floor,
 * Boston, MA  02110-1301, USA.
 *
 * Written by Luis Campo Giralte <luis.camp0.2009@gmail.com>
 *
 */
#include "NetbiosProtocol.h"
#include <iomanip>

namespace aiengine {

NetbiosProtocol::NetbiosProtocol():
	Protocol("Netbios", IPPROTO_UDP) {}

NetbiosProtocol::~NetbiosProtocol() {

	anomaly_.reset();
}

bool NetbiosProtocol::check(const Packet &packet){

	int length = packet.getLength();

	if (length >= header_size) {
		if ((packet.getSourcePort() == 137)or(packet.getDestinationPort() == 137)or
			(packet.getSourcePort() == 138)or(packet.getDestinationPort() == 138)) {
			++total_valid_packets_;
			return true;
		}
	}
	++total_invalid_packets_;
	return false;
}

void NetbiosProtocol::setDynamicAllocatedMemory(bool value) {

	info_cache_->setDynamicAllocatedMemory(value);
	name_cache_->setDynamicAllocatedMemory(value);
}

bool NetbiosProtocol::isDynamicAllocatedMemory() const {

	return info_cache_->isDynamicAllocatedMemory();
}

uint64_t NetbiosProtocol::getCurrentUseMemory() const {

	uint64_t mem = sizeof(NetbiosProtocol);

	mem += info_cache_->getCurrentUseMemory();
	mem += name_cache_->getCurrentUseMemory();

	return mem;
}

uint64_t NetbiosProtocol::getAllocatedMemory() const {

	uint64_t mem = sizeof(NetbiosProtocol);

        mem += info_cache_->getAllocatedMemory();
        mem += name_cache_->getAllocatedMemory();

        return mem;
}

uint64_t NetbiosProtocol::getTotalAllocatedMemory() const {

        uint64_t mem = getAllocatedMemory();

	mem += compute_memory_used_by_maps();

	return mem;
}

uint64_t NetbiosProtocol::compute_memory_used_by_maps() const {

	uint64_t bytes = 0;

	std::for_each (name_map_.begin(), name_map_.end(), [&bytes] (PairStringCacheHits const &f) {
		bytes += f.first.size();
	});
	return bytes;
}

uint32_t NetbiosProtocol::getTotalCacheMisses() const {

	uint32_t miss = 0;

	miss = info_cache_->getTotalFails();
	miss += name_cache_->getTotalFails();

	return miss;
}

void NetbiosProtocol::releaseCache() {

        if (FlowManagerPtr fm = flow_mng_.lock(); fm) {
                auto ft = fm->getFlowTable();

                std::ostringstream msg;
                msg << "Releasing " << name() << " cache";

                infoMessage(msg.str());

                uint64_t total_cache_bytes_released = compute_memory_used_by_maps();
                uint64_t total_bytes_released_by_flows = 0;
                uint64_t total_cache_save_bytes = 0;
                uint32_t release_flows = 0;
                uint32_t release_name = name_map_.size();

                for (auto &flow: ft) {
                        if (SharedPointer<NetbiosInfo> info = flow->getNetbiosInfo(); info) {
                                total_bytes_released_by_flows += sizeof(info);

                                flow->layer7info.reset();
                                ++release_flows;
                                info_cache_->release(info);
                        }
                }
                // Some entries can be still on the maps and needs to be
                // retrieve to their existing caches
                for (auto &entry: name_map_) {
			total_cache_save_bytes += entry.second.sc->size() * (entry.second.hits - 1);
                        releaseStringToCache(name_cache_, entry.second.sc);
		}
                name_map_.clear();

		data_time_ = boost::posix_time::microsec_clock::local_time();

                msg.str("");
                msg << "Release " << release_name << " netbios names, " << release_flows << " flows";
                computeMemoryUtilization(msg, total_cache_bytes_released, total_bytes_released_by_flows, total_cache_save_bytes);
		infoMessage(msg.str());
        }
}

void NetbiosProtocol::releaseFlowInfo(Flow *flow) {

	if (auto info = flow->getNetbiosInfo(); info)
		info_cache_->release(info);
}

void NetbiosProtocol::attach_netbios_name(NetbiosInfo *info, const boost::string_ref &name) {

        if (!info->netbios_name) {
                if (StringMap::iterator it = name_map_.find(name); it != name_map_.end()) {
                        ++(it->second).hits;
                        info->netbios_name = (it->second).sc;
		} else {
                        if (SharedPointer<StringCache> name_ptr = name_cache_->acquire(); name_ptr) {
                                name_ptr->name(name.data(), name.length());
                                info->netbios_name = name_ptr;
                                name_map_.insert(std::make_pair(name_ptr->name(), name_ptr));
                        }
                }
        }
}

void NetbiosProtocol::processFlow(Flow *flow) {

	CPUCycle cycles(&total_cpu_cycles_);
	int length = flow->packet->getLength();
	total_bytes_ +=  length;
	++total_packets_;
	current_flow_ = flow;

       	SharedPointer<NetbiosInfo> info = flow->getNetbiosInfo();
       	if (!info) {
               	if (info = info_cache_->acquire(); !info) {
			logFailCache(info_cache_->name(), flow);
			return;
		}
               	flow->layer7info = info;
       	}

	// A minimum header is :
	//   - header_size , simmilar to a dns.
	//   - 32 bytes of the netbios name.
	//   - 4 bytes of type and class.

	if (length <= header_size + 36) {
		++total_events_;
		// Malformed header packet
	        if (current_flow_->getPacketAnomaly() == PacketAnomalyType::NONE)
                       	current_flow_->setPacketAnomaly(PacketAnomalyType::NETBIOS_BOGUS_HEADER);

               	anomaly_->incAnomaly(PacketAnomalyType::NETBIOS_BOGUS_HEADER);
               	return;
	}

	setHeader(flow->packet->getPayload());
	int offset = 0;
	for (int i = 0; i < 32; i = i + 2) {
		uint8_t ptr1 = header_->data[i + 1];
		uint8_t ptr2 = header_->data[i + 2];

		if (ptr1 < 'A' or ptr1 > 'P' or ptr2 < 'A' or ptr2 > 'P')
			break;

		uint8_t value = ((ptr1 - 'A') << 4) + (ptr2 - 'A');

		if (value == 32) { // space that we dont want
			netbios_name_[offset] = 0x00;
			break;
		}

		if ((int)value <= 20) // skip the strange nb characters
			continue;

		netbios_name_[offset] = value;
		++offset;
	}

	if (offset > 0) { // There is something to attach
		boost::string_ref nb_name(reinterpret_cast<char*>(netbios_name_), offset);
		attach_netbios_name(info.get(), nb_name);
	}
}

void NetbiosProtocol::statistics(std::basic_ostream<char> &out, int level, int32_t limit) const {

	showStatisticsHeader(out, level);

	if ((level > 5)and(flow_forwarder_.lock()))
		flow_forwarder_.lock()->statistics(out);
	if (level > 3) {
		info_cache_->statistics(out);
		name_cache_->statistics(out);
		if (level > 4)
			name_map_.show(out, "\t",limit);
	}
}

void NetbiosProtocol::statistics(Json &out, int level) const {

	showStatisticsHeader(out, level);
}

void NetbiosProtocol::increaseAllocatedMemory(int value) {

        info_cache_->create(value);
        name_cache_->create(value);
}

void NetbiosProtocol::decreaseAllocatedMemory(int value) {

        info_cache_->destroy(value);
        name_cache_->destroy(value);
}

CounterMap NetbiosProtocol::getCounters() const {
	CounterMap cm;

        cm.addKeyValue("packets", total_packets_);
        cm.addKeyValue("bytes", total_bytes_);

        return cm;
}

#if defined(PYTHON_BINDING) || defined(RUBY_BINDING)
#if defined(PYTHON_BINDING)
boost::python::dict NetbiosProtocol::getCacheData(const std::string &name) const {
#elif defined(RUBY_BINDING)
VALUE NetbiosProtocol::getCacheData(const std::string &name) const {
#endif
        if (boost::iequals(name, "name"))
		return addMapToHash(name_map_);

	StringMap empty {"", ""};

        return addMapToHash(empty);
}

#if defined(PYTHON_BINDING)
SharedPointer<Cache<StringCache>> NetbiosProtocol::getCache(const std::string &name) {

        if (boost::iequals(name, "name"))
                return name_cache_;

        return nullptr;
}

#endif

#endif

void NetbiosProtocol::statistics(Json &out, const std::string &map_name, int32_t limit) const {

        if (boost::iequals(map_name, "names")) {
                name_map_.show(out, limit);
        }
}

void NetbiosProtocol::resetCounters() {

	reset();

	total_events_ = 0;
}

} // namespace aiengine
